import glob
import os
import random

from PIL import Image
from loguru import logger

CAPTCHA_WIDTH = 520
CAPTCHA_HEIGHT = 320


def generatr_xml(jpgPath, xmlPath, xmin, ymin, xmax, ymax):
    jpgName = os.path.split(jpgPath)[1]
    with open(xmlPath, 'w', encoding="utf-8") as xml_file:
        xml_file.write('<annotation>\n')
        xml_file.write('\t<folder>JPEGImages</folder>\n')
        xml_file.write('\t<filename>' + jpgName + '</filename>\n')
        xml_file.write('\t<path>' + jpgPath + '</path>\n')
        xml_file.write('\t<source>\n')
        xml_file.write('\t\t<database>' + 'Unknown' + '</database>\n')
        xml_file.write('\t</source>\n')
        xml_file.write('\t<size>\n')
        xml_file.write('\t\t<width>' + str(CAPTCHA_WIDTH) + '</width>\n')
        xml_file.write('\t\t<height>' + str(CAPTCHA_HEIGHT) + '</height>\n')
        xml_file.write('\t\t<depth>' + str(3) + '</depth>\n')
        xml_file.write('\t</size>\n')
        xml_file.write('\t<segmented>0</segmented>\n')

        xml_file.write('\t<object>\n')
        xml_file.write('\t\t<name>vcode</name>\n')
        xml_file.write('\t\t<pose>Unspecified</pose>\n')
        xml_file.write('\t\t<truncated>0</truncated>\n')
        xml_file.write('\t\t<difficult>0</difficult>\n')
        xml_file.write('\t\t<bndbox>\n')
        xml_file.write('\t\t\t<xmin>' + str(xmin) + '</xmin>\n')
        xml_file.write('\t\t\t<ymin>' + str(ymin) + '</ymin>\n')
        xml_file.write('\t\t\t<xmax>' + str(xmax) + '</xmax>\n')
        xml_file.write('\t\t\t<ymax>' + str(ymax) + '</ymax>\n')
        # xml_file.write('\t\t\t<angle>' + str(4) + '</angle>\n')
        xml_file.write('\t\t</bndbox>\n')
        # xml_file.write('\t\t<extra/>\n')
        xml_file.write('\t</object>\n')
        xml_file.write('</annotation>')


index = 0
for target in os.listdir("./static"):
    image_target = Image.open(os.path.join("./static", target))

    label_offset_x_base = 260
    label_offset_y_base = 120

    count = 0
    file_paths = glob.glob(f'outputs/*.jpg')

    for file_path in file_paths:
        offset_y = random.randint(100, 220)
        offset_x = random.randint(100, 420)

        image_path = os.path.join(file_path)
        image = Image.open(image_path)
        image.paste(image_target, (offset_x, offset_y), image_target)
        # # label = f'0 {label_offset_x / CAPTCHA_WIDTH} {label_offset_y / CAPTCHA_HEIGHT} 0.16596774 0.24170968'
        jpgPath = f'G:\\workspace\\python\\yolov8-all-in-one\\datasets\\VOCdevkit\\JPEGImages\\vcode_{index}_{count}.jpg'
        image.save(jpgPath)
        # logger.debug(f'generated captcha file captcha_{count}.png')
        # with open(f'G:\\workspace\\python\\yolov8-all-in-one\\datasets\VOCdevkit\\txt\\vcode_{index}_{count}.txt',
        #           'w') as f:
        #     f.write(label)
        xmlPath = f'G:\\workspace\\python\\yolov8-all-in-one\\datasets\\VOCdevkit\\Annotations\\vcode_{index}_{count}.xml'
        xmax = offset_x + image_target.size[0]
        ymax = offset_y + image_target.size[1]
        generatr_xml(jpgPath, xmlPath, offset_x, offset_y, xmax, ymax)
        logger.debug(f'generated label file captcha_{count}.xml')
        count += 1
    index += 1
